# deepseek_client.py
# 基于 DeepSeek API 的客户端
import requests
import json
from typing import Generator
import httpx

# 从 config.json 读取配置
with open("config.json", encoding="utf-8") as f:
    config = json.load(f)

DEEPSEEK_API_KEY = config["DEEPSEEK_API_KEY"]
DEEPSEEK_URL = config["DEEPSEEK_URL"]
MODEL_NAME = config["MODEL_NAME"]

def call_deepseek(prompt):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}"
    }

    payload = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "user", "content": prompt}
        ]
    }

    response = requests.post(DEEPSEEK_URL, headers=headers, json=payload)
    result = response.json()
    return result["choices"][0]["message"]["content"]

# 新增：流式推理生成器
def stream_deepseek(prompt) -> Generator[str, None, None]:
    """Yield generated content chunks from DeepSeek API using SSE."""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}"
    }
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "stream": True,
        "max_tokens": 2048,
        "stream_options": {"include_usage": True},
    }

    # 使用 stream=True 触发增量输出
    response = requests.post(DEEPSEEK_URL, headers=headers, json=payload, stream=True)
    # iter_lines 将保持连接并逐行读取
    for line in response.iter_lines(decode_unicode=True):
        if not line:
            continue
        # DeepSeek/SSE 行以 "data: " 开头
        if line.startswith("data: "):
            data_str = line[6:]
            # 结束标志
            if data_str.strip() == "[DONE]":
                break
            # 解析 JSON，获取增量内容
            try:
                data_json = json.loads(data_str)
                # usage 块 choices 为空，跳过
                choices = data_json.get("choices", [])
                if not choices:
                    continue
                delta = choices[0]["delta"].get("content", "")
                if delta:
                    yield delta
            except json.JSONDecodeError:
                # 忽略无法解析的片段
                continue

# 异步流式生成器
async def async_stream_deepseek(prompt):
    """异步生成器，从 DeepSeek API 生成内容（SSE）。"""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}"
    }
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "stream": True,
        "max_tokens": 2048,
        "stream_options": {"include_usage": True},
    }
    async with httpx.AsyncClient(timeout=None) as client:
        async with client.stream("POST", DEEPSEEK_URL, headers=headers, json=payload) as resp:
            async for line in resp.aiter_lines():
                if not line:
                    continue
                if line.startswith("data: "):
                    data_str = line[6:]
                    if data_str.strip() == "[DONE]":
                        break
                    try:
                        data_json = json.loads(data_str)
                        choices = data_json.get("choices", [])
                        if not choices:
                            continue
                        delta = choices[0]["delta"].get("content", "")
                        if delta:
                            yield delta
                    except json.JSONDecodeError:
                        continue
